import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Map;

import javax.xml.parsers.ParserConfigurationException;

import org.xml.sax.SAXException;

import com.beust.jcommander.JCommander;
import com.gq.utils.Crypto;


public class SQLiteEncrypt {
	
	public static boolean isDecrypt = false;
	public static String password;
	private static String dbPath;
	
	public static void encryptDatabase(SQLArg args) throws IOException {
		System.out.println("> start encrypt sqlite database");
		InputStream is = null;
		if (args.getSchemaFile() == null) {
			is = SQLiteEncrypt.class.getClassLoader().getResource("data.xml").openStream();
		} else {
			is = new FileInputStream(args.getSchemaFile());
		}
		Connection conn = null;
		Statement stm = null;
		ResultSet rs = null;
		SQLiteUtils sqliteutil = new SQLiteUtils(dbPath);
		try {
			Crypto c = new Crypto(password);
			ArrayList<Table> tables = DataParser.parse(is);
			Class.forName("org.sqlite.JDBC");
			conn = DriverManager.getConnection("jdbc:sqlite:"+dbPath);
			stm = conn.createStatement();
			
			for (Table table : tables) {
				System.out.println("---> encrypt table: "+table.getTableName());
				Map<String, String> columns = sqliteutil.getMetaData(table.getTableName());
				rs = stm.executeQuery("SELECT * FROM "+table.getTableName());
				while(rs.next()) {
					String sql = "UPDATE " + table.getTableName() + " SET ";
					ArrayList<Object> params = new ArrayList<Object>();
					System.out.println("---id: "+rs.getString(table.getPrimaryKey()));
					for (String col : table.getEncryptColumn()) {
						if (!columns.keySet().contains(col)) continue;
						String columnVaue = rs.getString(col);
						columnVaue = isDecrypt ? c.decryptFromBase64(columnVaue) : c.encryptToBase64(columnVaue);
						sql += col + "=?, ";
						params.add(columnVaue);
					}
					sql = sql.substring(0, sql.length() - 2);
					sql += String.format(" WHERE %s=?",table.getPrimaryKey());
					params.add(rs.getInt("id"));
					PreparedStatement ps = conn.prepareStatement(sql);
					for(int i=0;i<params.size();i++) {
						if (i == params.size() - 1) {
							ps.setInt(i+1, (Integer) params.get(i));
							continue;
						} 
						ps.setString(i+1, (String) params.get(i));
					}
					ps.execute();
				}
			}
			
			
		} catch (SAXException e) {
			e.printStackTrace();
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		} catch (SQLException e) {
			e.printStackTrace();
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			try {
				if (rs != null) rs.close();
				if (stm != null) stm.close();
				if (conn != null) conn.close();
			} catch(Exception e) {
			}
		}
		
		System.out.println("> Finished");
		
	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		
		SQLArg env = new SQLArg();
		JCommander cmd = new JCommander(env);
		String currentDir = System.getProperty("user.dir");
		try {
			cmd.parse(args);
			dbPath = env.getDbPath();
			if (!dbPath.matches("[a-zA-Z]:.+")) {
				dbPath = new File(currentDir, dbPath).toString();
			}
			File f = new File(dbPath);
			if (!f.isFile()) {
				System.out.println("the db file is not exists");
				System.exit(1);
			}
			isDecrypt = env.isDecrypt();
			password = env.getPwd();
			encryptDatabase(env);
			if (args.length == 0) throw new Exception("show usage");
		} catch(Exception e) {
			cmd.usage();
		}
	}

}
